import torch

# 在 PyTorch 中，unsqueeze 和 expand 是两个常用的操作，用于调整张量的维度和形状。

# unsqueeze(dim): 在指定的维度上增加一个维度。例如，tensor.unsqueeze(1) 会在第二个维度上添加一个维度。

# expand(*sizes): 扩展张量的形状，使其与指定的形状匹配。这个操作会自动地在维度上复制数据。

# expand 并不会真正复制张量的数据，而只是通过在存储中创建一个视图来实现形状的扩展。如果在扩展的过程中涉及到实际的复制，可以使用 repeat。

x = torch.tensor([1, 2, 3])
y = x.unsqueeze(0)  # 在第一个维度上添加一个维度
print(y.shape)  # 输出 torch.Size([1, 3])

import torch

a = torch.Tensor([[1], [2], [3],[4]])
# 未使用expand（）函数前的a
print('a.size: ', a.size())
print('a: ', a)

b = a.expand(4, 2)
# 使用expand（）函数后的输出
print('a.size: ', a.size())
print('a: ', a)
print('b.size: ', b.size())
print('b: ', b)

a = torch.Tensor([[[[1,2], [2,3], [3,4],[4,5]]]])   # [1,1,4,2]
b = a.expand(2, 1, 4, 2)
c = a.expand(1, 2, 4, 2)
# 使用expand（）函数后的输出
print('a.size: ', a.size())

print('b.size: ', b.size())
print('b: ', b)

print('c.size: ', c.size())
print('c: ', c)

# 由此可见，只要是单维度均可进行扩展，但是若非单维度会报错   下面语句会报错
# b2 = b.expand(4, 1, 4, 2)  # b: torch.Size([2, 1, 4, 2])
# print('b2.size: ', b2.size())


# 示例目标张量
target_tensor = torch.rand((2, 5, 5, 30))  # 假设 batch_size=2, S=5, 第四个维度有30个值

# 生成掩码
noo_mask = target_tensor[:, :, :, 4] == 0

# 打印结果
print(noo_mask.shape)  # 输出应该是 torch.Size([2, 5, 5])

noo_mask = noo_mask.unsqueeze(-1)

# 打印结果
print(noo_mask.shape)  # 输出应该是 torch.Size([2, 5, 5, 1])